package com.feidee.fd.sml.algorithm.component.feature

import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.feature.RFormula
import org.apache.spark.sql.DataFrame

/**
  * @Author: dongguosheng
  * @Date: 2019/3/20 19:43
  * @Review
  * @Email: guosheng_dong@sui.com
  */
case class RFormulaEncoderParam(
                                 override val input_pt: String,
                                 override val output_pt: String,
                                 override val hive_table: String,
                                 override val flow_time: String,
                                 // 无用参数
                                 override val inputCol: String,
                                 override val outputCol: String,
                                 override val preserveCols: String,
                                 override val modelPath: String,
                                 // 设置 R 语音 Model Formulae 公式，使用之前必须先调用这个函数设置参数例 如"y ~ x + z"
                                 formula: String
                               ) extends FeatureParam {

  def this() = this(null, null, null, null, "input", "features", null, null, null)

  override def verify(): Unit = {
    require(tool.isNotNull(input_pt), "param input_pt can't be null")
    require(tool.isNotNull(output_pt), "param output_pt can't be null")
    require(tool.isNotNull(outputCol), "输出字段 outputCol 不能为空")
    require(tool.isNotNull(formula), "param formula can't be null")
    require(!tool.isNotNull(preserveCols) || !preserveCols.contains(outputCol),
      "保留字段 preserveCols 不能与输出字段 outputCol 出现重合")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("formula" -> formula)
    map
  }
}


class RFormulaEncoder extends AbstractFeatureEncoder[RFormulaEncoderParam] {
  override def setUp(param: RFormulaEncoderParam, data: DataFrame): Array[PipelineStage] = {
    val rformula = new RFormula()
      .setFormula(param.formula)
      .setFeaturesCol(param.outputCol)  // 此处 featuresCol 即为处理结果列

    // P.S. labelCol 无需设置，可以由 formula 解析出来，但是后面也不做该列的保存，因此可以忽略掉该参数
    // （要保存标签列可通过 preserveCols 设置）

    Array(rformula)
  }
}

object RFormulaEncoder {

  def apply(paramStr: String): Unit = {
    new RFormulaEncoder()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    RFormulaEncoder(args(0))
  }
}